# Check if packages in requirements.txt are already satisfied in existing conda environments
# If not, create a new environment and install them

param (
    [string]$RequirementsPath = "requirementsFiles/pdf_ocr/requirements.txt",
    [string]$OutputDir = "conda_environments",
    [string]$NewEnvName = "pdf_ocr_env"
)

# Define output directory
if (-not (Test-Path $OutputDir)) {
    New-Item -ItemType Directory -Path $OutputDir | Out-Null
}

# Read requirements.txt file content
if (-not (Test-Path $RequirementsPath)) {
    Write-Host "Error: requirements.txt file not found: $RequirementsPath" -ForegroundColor Red
    exit 1
}

$requirements = Get-Content $RequirementsPath

# Extract main packages with their versions
$mainPackages = @{}
$pythonVersion = "3.9"  # Default Python version

foreach ($line in $requirements) {
    # Ignore comments and empty lines
    if ($line -match "^\s*#" -or $line -match "^\s*$") {
        continue
    }
    
    # Extract package name and version
    if ($line -match "^([a-zA-Z0-9_\-]+)(==|>=|<=|~=|!=|>|<)(.+)") {
        $packageName = $matches[1].ToLower()
        $operator = $matches[2]
        $version = $matches[3]
        
        # Check if it's a main package
        if ($packageName -eq "python") {
            $pythonVersion = $version
            $mainPackages[$packageName] = @{
                "operator" = $operator
                "version" = $version
            }
        }
        elseif ($packageName -eq "torch") {
            $mainPackages[$packageName] = @{
                "operator" = $operator
                "version" = $version
            }
        }
    }
    elseif ($line -match "^([a-zA-Z0-9_\-]+)") {
        $packageName = $matches[1].ToLower()
        
        # Check if it's a main package without version
        if ($packageName -eq "python" -or $packageName -eq "torch") {
            $mainPackages[$packageName] = @{
                "operator" = ""
                "version" = ""
            }
        }
    }
}

Write-Host "Found the following main packages in requirements.txt:"
foreach ($package in $mainPackages.Keys) {
    $versionInfo = $mainPackages[$package]
    if ($versionInfo.operator -and $versionInfo.version) {
        Write-Host "- $package $($versionInfo.operator) $($versionInfo.version)"
    } else {
        Write-Host "- $package (any version)"
    }
}

# Get all conda environments
$envs = conda env list | Where-Object { $_ -notmatch "^#" -and $_.Trim() -ne "" } | ForEach-Object {
    $parts = $_ -split '\s+', 2
    if ($parts.Count -ge 1 -and $parts[0] -ne "") {
        $parts[0]
    }
}

# Check if main packages are installed in each environment
$foundEnvironment = $false
$compatibleEnvs = @()

foreach ($env in $envs) {
    Write-Host "`nChecking environment: $env" -ForegroundColor Cyan
    
    # Get installed packages with versions in the environment
    $installedPackages = @{}
    conda list --name $env | Where-Object { $_ -notmatch "^#" -and $_.Trim() -ne "" } | ForEach-Object {
        $parts = $_ -split '\s+'
        if ($parts.Count -ge 3) {
            $packageName = $parts[0].ToLower()
            $version = $parts[1]
            $installedPackages[$packageName] = $version
        }
    }
    
    # Check if all main packages with correct versions are installed
    $missingOrWrongVersionPackages = @()
    
    foreach ($package in $mainPackages.Keys) {
        $versionInfo = $mainPackages[$package]
        
        if (-not $installedPackages.ContainsKey($package)) {
            $missingOrWrongVersionPackages += "$package (missing)"
            continue
        }
        
        $installedVersion = $installedPackages[$package]
        
        # If no specific version required, any version is fine
        if (-not $versionInfo.operator -or -not $versionInfo.version) {
            continue
        }
        
        # Compare versions based on operator
        $versionSatisfied = $false
        switch ($versionInfo.operator) {
            "==" { $versionSatisfied = [version]$installedVersion -eq [version]$versionInfo.version }
            ">=" { $versionSatisfied = [version]$installedVersion -ge [version]$versionInfo.version }
            "<=" { $versionSatisfied = [version]$installedVersion -le [version]$versionInfo.version }
            ">"  { $versionSatisfied = [version]$installedVersion -gt [version]$versionInfo.version }
            "<"  { $versionSatisfied = [version]$installedVersion -lt [version]$versionInfo.version }
            # For more complex operators like ~= or !=, we'd need more sophisticated version comparison
            default { $versionSatisfied = $true }  # Assume satisfied for unsupported operators
        }
        
        if (-not $versionSatisfied) {
            $missingOrWrongVersionPackages += "$package (found $installedVersion, required $($versionInfo.operator) $($versionInfo.version))"
        }
    }
    
    if ($missingOrWrongVersionPackages.Count -eq 0 -and $mainPackages.Count -gt 0) {
        Write-Host "Environment $env satisfies all main package requirements" -ForegroundColor Green
        Write-Host "Installed versions:"
        foreach ($package in $mainPackages.Keys) {
            if ($installedPackages.ContainsKey($package)) {
                Write-Host ("- {0}: {1}" -f $package, $installedPackages[$package])
            }
        }
        $compatibleEnvs += $env
        $foundEnvironment = $true
    } else {
        Write-Host "Environment $env has issues with the following packages:" -ForegroundColor Yellow
        $missingOrWrongVersionPackages | ForEach-Object { Write-Host "- $_" -ForegroundColor Yellow }
    }
}

# If no environment satisfies the requirements, create a new one
if (-not $foundEnvironment) {
    Write-Host "`nNo environment found that satisfies all main package requirements. Create a new environment? (Y/N)" -ForegroundColor Yellow
    $response = Read-Host
    
    if ($response -eq "Y" -or $response -eq "y") {
        Write-Host "Creating new environment: $NewEnvName" -ForegroundColor Cyan
        
        # Create new environment with Python
        Write-Host "Creating environment with Python $pythonVersion..."
        conda create --name $NewEnvName python=$pythonVersion -y
        
        # Install main packages with specific versions
        foreach ($package in $mainPackages.Keys) {
            if ($package -ne "python") {
                $versionInfo = $mainPackages[$package]
                if ($versionInfo.operator -and $versionInfo.version) {
                    Write-Host "Installing $package $($versionInfo.operator) $($versionInfo.version)..."
                    conda install --name $NewEnvName "$package$($versionInfo.operator)$($versionInfo.version)" -y
                } else {
                    Write-Host "Installing $package (latest version)..."
                    conda install --name $NewEnvName $package -y
                }
            }
        }
        
        # Install remaining packages
        Write-Host "Installing remaining packages with pip..."
        conda run --name $NewEnvName pip install -r $RequirementsPath
        
        Write-Host "`nEnvironment $NewEnvName has been created with required packages" -ForegroundColor Green
        
        # Export environment
        $fullOutputPath = Join-Path $OutputDir "${NewEnvName}_full.yml"
        $compatOutputPath = Join-Path $OutputDir "${NewEnvName}_compatible.yml"
        
        conda env export --name $NewEnvName > $fullOutputPath
        conda env export --name $NewEnvName --no-builds > $compatOutputPath
        
        Write-Host "Environment exported to: $fullOutputPath and $compatOutputPath"
    } else {
        Write-Host "Operation cancelled" -ForegroundColor Yellow
    }
} else {
    Write-Host "`nThe following environments satisfy all main package requirements:" -ForegroundColor Green
    $compatibleEnvs | ForEach-Object { Write-Host "- $_" }
    
    Write-Host "`nYou can activate these environments with the following commands:"
    $compatibleEnvs | ForEach-Object { Write-Host "conda activate $_" }
}

# Add this function at the beginning of your script
function Clean-YamlForCurrentPlatform {
    param (
        [string]$YamlPath,
        [string]$OutputPath
    )
    
    $content = Get-Content $YamlPath -Raw
    
    # Remove Linux-specific packages
    $linuxPackages = @(
        'readline', 'ncurses', 'libxcrypt', 'libuuid', 'libnsl', 
        'libgcc-ng', 'ld_impl_linux-64', 'libstdcxx-ng'
    )
    
    $yamlLines = $content -split "`n"
    $cleanedLines = @()
    $inDependencies = $false
    
    foreach ($line in $yamlLines) {
        if ($line -match "^dependencies:") {
            $inDependencies = $true
            $cleanedLines += $line
            continue
        }
        
        if ($inDependencies -and $line -match "^\s*-\s+(.+?)=") {
            $packageName = $matches[1].Trim()
            $isLinuxPackage = $false
            
            foreach ($linuxPkg in $linuxPackages) {
                if ($packageName -eq $linuxPkg) {
                    $isLinuxPackage = $true
                    break
                }
            }
            
            if (-not $isLinuxPackage) {
                $cleanedLines += $line
            }
        } else {
            $cleanedLines += $line
        }
    }
    
    $cleanedContent = $cleanedLines -join "`n"
    Set-Content -Path $OutputPath -Value $cleanedContent
    
    Write-Host "Created platform-compatible YAML at: $OutputPath" -ForegroundColor Green
} 